use std::path::Path;

use ash::vk;

use crate::{vulkan::{vk_utils::{utils_info, utils_pipeline}, VkContext}, VK_CHECK_UNSAFE, VK_LOGE, VK_LOGP, system::render::{self, render_pass}};

use super::{VkDescriptorSetLayoutGenerator, VkPipelineGenerator, VkPipelineState, VkTexture, VkWriteDescriptorSetGenerator, VkPipeline};

impl VkPipelineState {
    pub fn default() -> Self {
        let vertex_input_state          = utils_info::get_pipeline_vertex_input_state_ci();
        let input_assembly_state        = utils_info::get_pipeline_input_assembly_state_ci();
        let tessellation_state          = utils_info::get_pipeline_tessellation_state_ci();
        let viewport_state              = utils_info::get_pipeline_viewport_state_ci();
        let rasterization_state         = utils_info::get_pipeline_rasterization_state_ci();
        let multisample_state           = utils_info::get_pipeline_multisample_state_ci();
        let depth_stencil_state         = utils_info::get_pipeline_depth_stencil_state_ci();
        let color_blend_state           = utils_info::get_pipeline_color_blend_state_ci();
        let dynamic_state               = utils_info::get_pipeline_dynamic_state_ci();

        Self {
            vertex_input_state          : vertex_input_state,
            input_assembly_state        : input_assembly_state,
            tessellation_state          : tessellation_state,
            viewport_state              : viewport_state,
            rasterization_state         : rasterization_state,
            multisample_state           : multisample_state,
            depth_stencil_state         : depth_stencil_state,
            color_blend_state           : color_blend_state,
            dynamic_state               : dynamic_state
        }
    }
}

impl VkPipelineGenerator {
    pub fn new(descriptor_set_layout_count : usize) -> Self {
        let pipeline_state              = VkPipelineState::default();

        let mut descriptor_set_layout_generators = vec![];
        for index in 0..descriptor_set_layout_count as usize {
            descriptor_set_layout_generators.push(VkDescriptorSetLayoutGenerator::new());
        }

        Self {
            state                               : pipeline_state,
            descriptor_set_layout_count         : descriptor_set_layout_count,
            descriptor_set_layout_generators    : descriptor_set_layout_generators,
            pipeline_shader_stage_cis           : vec![]
        }
    }

    pub fn add_shader(&mut self, vk_context: &VkContext, path: &Path, stage: vk::ShaderStageFlags, entry: *const i8) {
        let mut shader_stage_ci = utils_info::get_pipeline_shader_stage_ci();
        shader_stage_ci.stage       = stage;
        shader_stage_ci.module      = utils_pipeline::load_module(vk_context, path);
        shader_stage_ci.p_name      = entry;

        self.pipeline_shader_stage_cis.push(shader_stage_ci);
    }

    pub fn add_set_binding(&mut self,
        set                     : usize,
        binding                 : u32,                      // 对应Shader中的Binding号
        descriptor_type         : vk::DescriptorType,
        descriptor_count        : u32,
        stage_flags             : vk::ShaderStageFlags,
        p_immutable_samplers    : *const vk::Sampler        // 暂时不用
    ) {
        if set >= self.descriptor_set_layout_count {
            VK_LOGE!("set[{set}] >= self.descriptor_set_layout_count[{}]", self.descriptor_set_layout_count);
            return
        }
        self.descriptor_set_layout_generators[set].add_binding(binding, descriptor_type, descriptor_count, stage_flags, p_immutable_samplers)
    }

    pub fn generate(&mut self, vk_context: &VkContext, render_pass: &vk::RenderPass, subpass: u32) -> VkPipeline {
        VkPipeline::new(vk_context, self, render_pass, subpass)
    }
}

impl VkPipeline {
    pub fn new(vk_context: &VkContext, pipeline_generator: &mut VkPipelineGenerator, render_pass: &vk::RenderPass, subpass: u32) -> Self {
        // Link Pipeline CreateInfo
        let mut pipeline_ci = utils_info::get_pipeline_ci();
        pipeline_ci.p_vertex_input_state        = &pipeline_generator.state.vertex_input_state;
        pipeline_ci.p_input_assembly_state      = &pipeline_generator.state.input_assembly_state;
        pipeline_ci.p_tessellation_state        = &pipeline_generator.state.tessellation_state;
        pipeline_ci.p_viewport_state            = &pipeline_generator.state.viewport_state;
        pipeline_ci.p_rasterization_state       = &pipeline_generator.state.rasterization_state;
        pipeline_ci.p_multisample_state         = &pipeline_generator.state.multisample_state;
        pipeline_ci.p_depth_stencil_state       = &pipeline_generator.state.depth_stencil_state;
        pipeline_ci.p_color_blend_state         = &pipeline_generator.state.color_blend_state;
        pipeline_ci.p_dynamic_state             = &pipeline_generator.state.dynamic_state;

        pipeline_ci.render_pass                 = *render_pass;
        pipeline_ci.subpass                     = subpass;

        pipeline_ci.stage_count                 = pipeline_generator.pipeline_shader_stage_cis.len() as u32;
        pipeline_ci.p_stages                    = pipeline_generator.pipeline_shader_stage_cis.as_ptr();
        // Create & Link DescriptorSetLayout belonging to Pipeline
        let mut descriptor_set_layouts = vec![];
        for descriptor_set_layout_generator in &pipeline_generator.descriptor_set_layout_generators {
            descriptor_set_layouts.push(descriptor_set_layout_generator.generate(vk_context))
        }
        let mut pipeline_layout_ci = utils_info::get_pipeline_layout_ci();
        pipeline_layout_ci.set_layout_count     = descriptor_set_layouts.len() as u32;
        pipeline_layout_ci.p_set_layouts        = descriptor_set_layouts.as_ptr();
        let pipeline_layout = match VK_CHECK_UNSAFE!( vk_context.vk_device.create_pipeline_layout(&pipeline_layout_ci, None)) {
            Some(inner_value)   => inner_value,
            None                => VK_LOGP!("Failed to create_pipeline_layout")
        };
        pipeline_ci.layout = pipeline_layout;
        // Create Pipeline
        let pipeline_handler = match unsafe{ vk_context.vk_device.create_graphics_pipelines(vk::PipelineCache::null(), &[pipeline_ci], None) } {
            Ok(pipeline_handlers) => pipeline_handlers[0],
            Err((pipeline_handlers, error)) => {
                VK_LOGE!("{:?}", error);
                pipeline_handlers[0]
            }
        };
        
        Self {
            handler                 : pipeline_handler,
            layout                  : pipeline_layout,
            descriptor_set_layouts  : descriptor_set_layouts
        }
    }
}

impl VkDescriptorSetLayoutGenerator {
    pub fn new() -> Self {
        Self {
            layout_bindings: vec![]
        }
    }

    /// 增加Descriptor的Layout到预加载中
    pub fn add_binding(&mut self, 
        binding                 : u32,                      // 对应Shader中的Binding号
        descriptor_type         : vk::DescriptorType,
        descriptor_count        : u32,
        stage_flags             : vk::ShaderStageFlags,
        p_immutable_samplers    : *const vk::Sampler        // 暂时不用
    ) {
        let mut descriptor_set_layout_binding = utils_info::get_descriptor_set_layout_binding();
        descriptor_set_layout_binding.binding               = binding;
        descriptor_set_layout_binding.descriptor_type       = descriptor_type;
        descriptor_set_layout_binding.descriptor_count      = descriptor_count;
        descriptor_set_layout_binding.stage_flags           = stage_flags;

        self.layout_bindings.push(descriptor_set_layout_binding);
    }

    /// 生成整个Binding对应的Descriptor Layout/Pool/Set
    /// layout也就是一个Descriptor Set的布局
    /// pool里面有多个layout对象
    pub fn generate(&self, vk_context: &VkContext) -> vk::DescriptorSetLayout {
        // 生成Layout
        let mut descriptor_set_layout_ci = utils_info::get_descriptor_set_layout_create_info();
        descriptor_set_layout_ci.binding_count      = self.layout_bindings.len() as u32;
        descriptor_set_layout_ci.p_bindings         = self.layout_bindings.as_ptr();

        match VK_CHECK_UNSAFE!(vk_context.vk_device.create_descriptor_set_layout(&descriptor_set_layout_ci, None)) {
            Some(inner_value)   => inner_value,
            None                => VK_LOGP!("Failed to create_descriptor_set_layout")
        }
    }

}

impl VkWriteDescriptorSetGenerator {
    pub fn new() -> Self {
        Self {
            write_descriptor_sets: vec![],

            descriptor_buffer_infos : vec![],
            descriptor_image_infos : vec![]
        }
    }

    pub fn add_buffer_info(&mut self, buffer: &vk::Buffer, descriptor_set: &vk::DescriptorSet, dst_binding: u32, range: u64, descriptor_buffer_info: vk::DescriptorBufferInfo) {
        // let descriptor_buffer_info = vk::DescriptorBufferInfo {
        //     buffer          : *buffer,
        //     offset          : 0,
        //     range           : range
        // };

        let write_descriptor_set = vk::WriteDescriptorSet {
            s_type                  : vk::StructureType::WRITE_DESCRIPTOR_SET,
            p_next                  : std::ptr::null(),
            dst_set                 : *descriptor_set,
            dst_binding             : dst_binding,
            dst_array_element       : 0,
            descriptor_count        : 1,
            descriptor_type         : vk::DescriptorType::UNIFORM_BUFFER,
            p_buffer_info           : &descriptor_buffer_info,
            p_image_info            : std::ptr::null(),
            p_texel_buffer_view     : std::ptr::null()
        };

        self.descriptor_buffer_infos.push(descriptor_buffer_info);
        self.write_descriptor_sets.push(write_descriptor_set);
    }

    pub fn add_image_info(&mut self, texture: &VkTexture, descriptor_set: &vk::DescriptorSet, dst_binding: u32, descriptor_image_info: &vk::DescriptorImageInfo, is_inner: bool) {
        // let descriptor_image_info = vk::DescriptorImageInfo {
        //     sampler         : texture.description.sampler,
        //     image_view      : texture.description.image_view,
        //     image_layout    : vk::ImageLayout::GENERAL
        // };

        let write_descriptor_set = vk::WriteDescriptorSet {
            s_type                  : vk::StructureType::WRITE_DESCRIPTOR_SET,
            p_next                  : std::ptr::null(),
            dst_set                 : *descriptor_set,
            dst_binding             : dst_binding,
            dst_array_element       : 0,
            descriptor_count        : 1,
            descriptor_type         : if is_inner { vk::DescriptorType::INPUT_ATTACHMENT } else { vk::DescriptorType::COMBINED_IMAGE_SAMPLER },
            p_buffer_info           : std::ptr::null(),
            p_image_info            : descriptor_image_info,
            p_texel_buffer_view     : std::ptr::null()
        };

        self.write_descriptor_sets.push(write_descriptor_set);
    }

    pub fn generate(&self) -> &Vec<vk::WriteDescriptorSet> {
        &self.write_descriptor_sets
    }

}